from enum import Enum
from typing import TypeVar, Optional, Generic

from pydantic.generics import GenericModel


class ErrorCode(Enum):
    """API错误码"""
    SUCCESS = (200, "请求成功")
    AI_DEFAULT_EXCEPTION = (90000, "系统异常")
    AI_ACCOUNT_ERROR = (90001, "账号异常")
    AI_RATE_LIMIT_ERROR = (90002, "账号限流")

    def __init__(self, code, msg):
        self.code = code
        self.message = msg


DataT = TypeVar("DataT")


class ApiResponse(GenericModel, Generic[DataT]):
    """
       通用返回数据
    """
    code: Optional[int] = None
    message: Optional[str] = None
    data: Optional[DataT] = None


def success(data: Optional[DataT] = None):
    return ApiResponse(code=ErrorCode.SUCCESS.code, message=ErrorCode.SUCCESS.message, data=data)


def error(code: Optional[int] = None, msg: Optional[str] = None):
    return ApiResponse(code=code, message=msg)


def custom_error(error_code: Optional[ErrorCode]):
    return ApiResponse(code=error_code.code, message=error_code.message)
